import base64
import copy
import os
from io import BytesIO

from PIL import Image
from openai import OpenAI
from qwen_vl_utils import process_vision_info
from transformers import AutoProcessor


class Qwen2VLModel:
    def __init__(self, model_name, detail="auto", system_prompt=None, system_message=None, base_url=None, api_key=None,
                 min_pixels=256 * 28 * 28, max_pixels=1280 * 28 * 28):
        # detail: auto, low, high
        self.base_url = base_url or os.getenv("OPENAI_BASE_URL")
        self.api_key = api_key or os.getenv("OPENAI_API_KEY")
        self.client = OpenAI(api_key=self.api_key, base_url=self.base_url)
        self.model_name = model_name
        self.detail = detail
        self.processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct")
        self.min_pixels = min_pixels
        self.max_pixels = max_pixels

        if system_prompt:
            self.system_prompt = system_prompt
        else:
            self.system_prompt = "You are a helpful assistant."
        if system_message:
            self.system_message = system_message
        else:
            self.system_message = ""

    def describe_image(self, image_path, query):
        image = Image.open(image_path)
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        img_base64 = base64.b64encode(buffered.getvalue()).decode()
        completion = self.client.chat.completions.create(
            model=self.model_name,
            messages=[
                {
                    "role": "system",
                    "content": self.system_prompt,
                },
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": query
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/png;base64,{img_base64}"
                            }
                        }
                    ]
                }
            ]
        )
        return completion.choices[0].message.content

    def chat_with_message(self, messages):
        messages_copy = copy.deepcopy(messages)
        for msg in messages_copy:
            if msg["role"] == "user":
                for i, content in enumerate(msg["content"]):
                    if content["type"] == "image":
                        processed_image = self.process_vision_info(content["image"])
                        msg["content"][i] = {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/png;base64,{processed_image}"
                            }
                        }

        completion = self.client.chat.completions.create(
            model=self.model_name,
            messages=messages_copy,
        )
        return completion.choices[0].message.content

    def chat(self, messages):
        messages_copy = copy.deepcopy(messages)
        for msg in messages_copy:
            if msg["role"] == "user":
                for i, content in enumerate(msg["content"]):
                    if content["type"] == "image":
                        processed_image = self.process_vision_info(content["image"])
                        msg["content"][i] = {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/png;base64,{processed_image}"
                            }
                        }

        completion = self.client.chat.completions.create(
            model=self.model_name,
            messages=messages_copy
        )
        return completion.choices[0].message.content

    def process_vision_info(self, image_file):
        messages = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "image",
                        "image": f"file://{image_file}",
                        "min_pixels": self.min_pixels,
                        "max_pixels": self.max_pixels,
                    }
                ]
            }
        ]
        image_inputs, _ = process_vision_info(messages)
        buffered = BytesIO()
        image_inputs[0].save(buffered, format="PNG")
        img_base64 = base64.b64encode(buffered.getvalue()).decode()
        return img_base64
